import numpy as np
import pandas as pd
from qtorch.strategy import Strategy

class TurtleTradingStrategy(Strategy):
    """海龟交易策略（20日突破）"""
    def __init__(self, entry_window=20, exit_window=10, atr_window=14):
        self.entry_window = entry_window
        self.exit_window = exit_window
        self.atr_window = atr_window
        
    def generate_signals(self, prices):
        highs = prices['high'] if isinstance(prices, pd.DataFrame) else prices
        lows = prices['low'] if isinstance(prices, pd.DataFrame) else prices
        closes = prices['close'] if isinstance(prices, pd.DataFrame) else prices
        
        # 计算突破通道
        entry_high = highs.rolling(window=self.entry_window).max()
        entry_low = lows.rolling(window=self.entry_window).min()
        exit_high = highs.rolling(window=self.exit_window).max()
        exit_low = lows.rolling(window=self.exit_window).min()
        
        # 计算ATR
        tr1 = highs - lows
        tr2 = abs(highs - closes.shift())
        tr3 = abs(lows - closes.shift())
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        atr = tr.rolling(window=self.atr_window).mean()
        
        # 生成信号
        long_entry = closes > entry_high
        short_entry = closes < entry_low
        long_exit = closes < exit_low
        short_exit = closes > exit_high
        
        signals = np.zeros(len(closes))
        position = 0
        
        for i in range(1, len(signals)):
            if position <= 0 and long_entry.iloc[i]:
                signals[i] = 1
                position = 1
            elif position >= 0 and short_entry.iloc[i]:
                signals[i] = -1
                position = -1
            elif position > 0 and long_exit.iloc[i]:
                signals[i] = 0
                position = 0
            elif position < 0 and short_exit.iloc[i]:
                signals[i] = 0
                position = 0
            else:
                signals[i] = signals[i-1]
                
        return signals.astype(int)